from ezdlops import BaseTrainer, to_persistent
import torch
from torch import nn
from torch.nn import functional as F
from ave2flame.networks import AudioEncoder
from tqdm import tqdm

class FlameDecoder(nn.Module):
  
  _expr_shape = 100
  _n_joints = 1
  
  @staticmethod
  def angles_to_unit_vector(phi, theta):
    """
    将方位角和俯仰角转换为单位向量。
    """
    x = torch.cos(theta) * torch.cos(phi)
    y = torch.cos(theta) * torch.sin(phi)
    z = torch.sin(theta)
    return torch.stack([x, y, z], dim=-1)
  
  def __init__(self, last_activate=True):
    super().__init__()
    self.shape_encode = nn.Sequential(
        nn.Linear(300, 128),
        nn.LeakyReLU(0.02, True),
        nn.Linear(128, 64),
        nn.Softplus()
    )
    
    input_dim = 64 + 512
    
    self.expr_mu = nn.Sequential(
        nn.Linear(input_dim, 256),
        nn.LeakyReLU(0.02, True),
        nn.Linear(256, self._expr_shape)
    )
    
    seq = [
        nn.Linear(input_dim, 256),
        nn.LeakyReLU(0.02, True),
        nn.Linear(256, self._expr_shape),
    ]
    if last_activate:
      seq.append(nn.Softplus())
    self.expr_var = nn.Sequential(*seq)
    
    self.pose_angle = nn.Sequential(
        nn.Linear(input_dim, 256),
        nn.LeakyReLU(0.02, True), 
        nn.Linear(256, 4 * self._n_joints), # joints * 2(phi, theta) * 2(mu, sigma)
    )
    
    seq = [
        nn.Linear(input_dim, 256),
        nn.LeakyReLU(0.02, True), 
        nn.Linear(256, 2 * self._n_joints), # 1(jaw, neck, 2eyes) * 1(angle) * 2(mu, sigma)
    ]
    if last_activate:
      seq.append(nn.Sigmoid()), # [0, 1] mu range [0, pi])
    self.rotate_angle = nn.Sequential(*seq)
    
  def forward(self, aud_feat, shape_feat):
    shape_enc = self.shape_encode(shape_feat)
    latent_feat = torch.cat([aud_feat, shape_enc], dim=-1)
    
    expr_var = self.expr_var(latent_feat)
    expr_mu = self.expr_mu(latent_feat)
    
    pose_angle = self.pose_angle(latent_feat)
    pose_angle_mu = torch.tanh(pose_angle[:, :2*self._n_joints]) * torch.pi
    pose_angle_var = torch.sigmoid(pose_angle[:, 2*self._n_joints:])
    
    rotate_angle = self.rotate_angle(latent_feat)
    rotate_mu = rotate_angle[:, :self._n_joints] * torch.pi # range [0, pi]
    rotate_var = rotate_angle[:, self._n_joints:] # range [0, 1]
  
    return expr_mu, expr_var, pose_angle_mu, pose_angle_var, rotate_mu, rotate_var
  
  @staticmethod
  def angles_to_unit_vector(phi, theta):
    x = torch.cos(theta) * torch.cos(phi)
    y = torch.cos(theta) * torch.sin(phi)
    z = torch.sin(theta)
    return torch.stack([x, y, z], dim=-1)
  
  def construct_z(self, expr_mu, expr_var, pose_angle_mu, pose_angle_var, rotate_mu, rotate_var):
    B = expr_mu.size(0)
    phi, theta = pose_angle_mu.view(B, 2, 1).unbind(1) # (B, 8) 4*mu_phi, 4*mu_theta -> (B, 2, 4) -> (B, 4), (B, 4)
    v = self.angles_to_unit_vector(phi, theta) # (B, 4), (B, 4) stack -> (B, 4, 3) mu_v
    rot_angle = rotate_mu.view(B, 1, 1) # (B, 4, 1) mu_angle
    pos  = (v * rot_angle).view(B, -1) # (B, 4, 3) mu_v * mu_angle -> (B, 4, 3) mu_pos -> (B, 12) mu_pos
    ret = torch.concat([expr_mu, pos], dim=-1) # (B, 100) mu_expr + (B, 12) mu_pos -> (B, 112) mu
    return ret
  
  def to_flame_feat(self, ave_feat, shape_feat):
    expr_mu, expr_var, pose_angle_mu, pose_angle_var, rotate_mu, rotate_var = self.forward(ave_feat, shape_feat)
    return self.construct_z(expr_mu, expr_var, pose_angle_mu, pose_angle_var, rotate_mu, rotate_var)
    
class FlameDiffDecoder(FlameDecoder):
  def __init__(self):
    super().__init__(last_activate = False)
    
  def forward(self, aud_feat, shape_feat):
    shape_enc = self.shape_encode(shape_feat)
    latent_feat = torch.cat([aud_feat, shape_enc], dim=-1)
    
    expr_mu = torch.tanh(self.expr_mu(latent_feat))
    expr_var = torch.sigmoid(self.expr_var(latent_feat))
    
    pose_angle = self.pose_angle(latent_feat)
    pose_angle_mu = torch.tanh(pose_angle[:, :2*self._n_joints])
    pose_angle_var = torch.sigmoid(pose_angle[:, 2*self._n_joints:])
    
    rotate_angle = self.rotate_angle(latent_feat)
    rotate_mu = torch.tanh(rotate_angle[:, :self._n_joints]) # range [-1, 1]
    rotate_var = torch.sigmoid(rotate_angle[:, self._n_joints:]) # range [0, 1]
  
    return expr_mu, expr_var, pose_angle_mu, pose_angle_var, rotate_mu, rotate_var
    

@to_persistent
class AveFace2Flame_Diff(BaseTrainer):
  
  _expr_shape = 100
  
  def __init__(self, 
               ave_path: str=None, 
               device: str='cpu',
               ):
    super().__init__()
    self.ave = AudioEncoder().to(device).eval()
    if not ave_path is None:
      ckpt = torch.load(ave_path)
      self.ave.load_state_dict({f'audio_encoder.{k}': v for k, v in ckpt.items()})
    
    self.decoder = FlameDecoder().to(device)
    self.decoder_diff = FlameDiffDecoder().to(device)
    self.device = device
    self.max_epoch = 300
    
    self.to(device)
    self.set_dir()
    
  @staticmethod
  def unit_vector_to_angles(v):
    """
    将单位向量转换为方位角和俯仰角。

    Args:
        unit_vector (torch.Tensor): 形状为 (..., 3) 的单位向量。

    Returns:
        tuple: 包含方位角 (phi) 和俯仰角 (theta) 的元组，均为 torch.Tensor。
    """
    # 确保输入是单位向量
    rot_angle = torch.norm(v, dim=-1, keepdim=True)
    unit_vector = v / rot_angle

    z = unit_vector[..., 2]
    xy_norm = torch.norm(unit_vector[..., :2], dim=-1)

    phi = torch.atan2(unit_vector[..., 1], unit_vector[..., 0])
    theta = torch.atan2(z, xy_norm)

    return phi, theta
  
  def loss(self, y, mu, var):
    """
    half gaussian nll
    """
    var = var + torch.finfo(var.dtype).eps
    loss = torch.log(var) + (y - mu)**2 / var
    return torch.mean(loss)  # 返回批次的平均损失
  
  def forward(self, aud, shape_feat, target_cur, target_prev):
    # gt
    with torch.no_grad():
      expr = target_cur[:, :self._expr_shape]
      pose = target_cur[:, self._expr_shape:].view(-1, 1, 3) # (B, 1 joint, 3 axis)
      rotate = pose.norm(dim=-1, keepdim=True) # (B, 1 rotate angle, 1)
      v = pose / rotate
      phi, theta = self.unit_vector_to_angles(v) # (B, 1 joint, 3 axis) -> (B, 1 joint) (phi, theta)
      
      expr_pre = target_prev[:, :self._expr_shape]
      pose_pre = target_prev[:, self._expr_shape:].view(-1, 1, 3) # (B, 1 joint, 3 axis)
      rotate_pre = pose_pre.norm(dim=-1, keepdim=True) # (B, 1 rotate angle, 1)
      v_pre = pose_pre / rotate_pre
      phi_pre, theta_pre = self.unit_vector_to_angles(v_pre) # (B, 1 joint, 3 axis) -> (B, 1 joint) (phi, theta)
      
      angle = torch.concat([phi, theta], dim=-1)
      ygt = torch.concat([expr, angle, rotate.squeeze(-1)], dim=-1)
      
      angle_pre = torch.concat([phi_pre, theta_pre], dim=-1)
      ygt_pre = torch.concat([expr_pre, angle_pre, rotate_pre.squeeze(-1)], dim=-1)
      ygt_diff = ygt - ygt_pre
     # pred current frame
      aud_feat = self.ave(aud)
      
    expr_mu, expr_var, pose_angle_mu, pose_angle_var, rotate_mu, rotate_var = self.decoder(aud_feat, shape_feat)
    mu = torch.concat([expr_mu, pose_angle_mu, rotate_mu], dim=-1)
    var = torch.concat([expr_var, pose_angle_var, rotate_var], dim=-1)
    # pred change
    expr_mu_diff, expr_var_diff, pose_angle_mu_diff, pose_angle_var_diff, rotate_mu_diff, rotate_var_diff = self.decoder_diff(aud_feat, shape_feat)
    mu_diff = torch.concat([expr_mu_diff, pose_angle_mu_diff, rotate_mu_diff], dim=-1)
    var_diff = torch.concat([expr_var_diff, pose_angle_var_diff, rotate_var_diff], dim=-1)
    mu_pre = mu - mu_diff
    var_pre = var_diff + var # 认为相互独立, 则方差相加
    
    # reconstruct
    pose_re = self.decoder.to_flame_feat(aud_feat, shape_feat)
    # loss
    loss_nll_cur = self.loss(ygt, mu, var)
    loss_nll_pre = self.loss(ygt_pre, mu_pre, var_pre)
    loss_nll_diff = self.loss(ygt_diff, mu_diff, var_diff)
    loss_l1 = (pose_re - target_cur).abs().mean()
    return loss_nll_cur, loss_nll_pre, loss_nll_diff, loss_l1
  
  def fit(self, dataloader, optimizer):
    for _ in range(self.max_epoch):
      dataiter = dataloader.__iter__()
      progress = tqdm(dataiter)
      loss_epoch = 0.
      for idx, data in enumerate(progress):
          aud_feat = data[0]
          
          flame_params, flame_params_pre, shape_feat = data[1]
          aud_feat = aud_feat.to(self.device)
          shape_feat = shape_feat.to(self.device)
          flame_params = flame_params.to(self.device)
          flame_params_pre = flame_params_pre.to(self.device)
          
          loss_nll_cur, loss_nll_pre, loss_nll_diff, loss_l1 = self(aud_feat, shape_feat, flame_params, flame_params_pre)
          loss = 0.45 * loss_nll_cur + 0.25 * loss_nll_pre + 0.25 * loss_nll_diff + 0.05 * loss_l1
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
          loss_epoch += loss.item()
          progress.set_description(f'epoch: {self._epoch}, loss: {loss.item():.3f}, loss_avg: {loss_epoch / (idx + 1):.3f}')   
          
          if self._iteration % 100 == 0:
            self._tb_writer.add_scalar('loss/loss_nll_cur', loss_nll_cur.item(), self._iteration)
            self._tb_writer.add_scalar('loss/loss_nll_pre', loss_nll_pre.item(), self._iteration)
            self._tb_writer.add_scalar('loss/loss_nll_diff', loss_nll_diff.item(), self._iteration)
            self._tb_writer.add_scalar('loss/loss_l1', loss_l1.item(), self._iteration)
          
          self._iteration += 1
          
      self._epoch += 1
      self._tb_writer.add_scalar('loss/loss_epoch', loss_epoch / idx, self._iteration)
      if self._epoch % 20 == 0:  
        self.save_auto()  